import pandas as pd

# import datasets

node_path = './CyberAssetGraphData-main/Dataset/Node.csv'
edge_path = './CyberAssetGraphData-main/Dataset/Link.csv'
node_save_path = './CyberAssetGraphData-main/Dataset/Cleaned_Node.csv'
edge_save_path = './CyberAssetGraphData-main/Dataset/Cleaned_Link.csv'

node_file = pd.read_csv(node_path, header=0)
edge_file = pd.read_csv(edge_path, header=0)

# add a unique identifier 'ID_num'
node_file['ID_num'] = range(1, len(node_file) + 1)

# replace ids with ID_num in Link dataset
node_file_dict = dict(zip(node_file['id'], node_file['ID_num']))
edge_file['source'] = edge_file['source'].map(node_file_dict)
edge_file['target'] = edge_file['target'].map(node_file_dict)

# delete the node of the error message
node_file = node_file.loc[~((node_file['type'] == 'Domain') & (node_file['industry'] == '[]'))]
# node_file.drop(((node_file['type'] == 'Domain') and (node_file['industry'] == '[]')))

# modify field error
for index, row in node_file.iterrows():
    if row['industry'] != '[]':
        s = ''
        for j in range(0, len(row['industry'])):
            if 'A' <= row['industry'][j] <= 'Z':
                s += row['industry'][j]
        if s != '':
            s = set(s)
        else:
            s = {}
        node_file.at[index, 'industry'] = s

# save cleaned datasets
node_file.to_csv(node_save_path, index=False)
edge_file.to_csv(edge_save_path, index=False)

print('save successfully')
